# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


def plot_balance(df, csv_root):
  plt.figure(figsize=(20, 20))
  df['wallet_balance'].plot(legend=True)
  plt.savefig(os.path.join(csv_root, 'balance_plot.png'))


def plot_position(df, csv_root):
  plt.figure(figsize=(20, 20))
  df['net_position'].plot(legend=True)
  plt.savefig(os.path.join(csv_root, 'position_plot.png'))


def plot_order_event(df, csv_root, exchange):
  plt.figure(figsize=(40, 40))
  if exchange == 'Bitmex':
    df.groupby('internal_order_id')['price'].plot(legend=False)
    accept, = plt.plot(df.loc[df['order_event_type'].isin([2, 7])]['price'],
                       color='black',
                       linestyle='None',
                       marker='o',
                       markersize=5)
    fill, = plt.plot(df.loc[df['order_event_type'].isin([4])]['fill_price'],
                     color='red',
                     linestyle='None',
                     marker='s',
                     markersize=10)
    cancel, = plt.plot(df.loc[df['order_event_type'].isin([14])]['price'],
                       color='blue',
                       linestyle='None',
                       marker='x',
                       markersize=10)
    plt.legend([accept, fill, cancel], ['submit & amend accept', 'fill', 'cancel'])
    plt.savefig(os.path.join(csv_root, 'order_event_plot.png'))
  elif exchange == 'Okex':
    line, = plt.plot(df[df['price'].notnull()]['price'], color='yellow')
    accept, = plt.plot(df.loc[df['order_event_type'].isin([2, 7])]['price'],
                       color='black',
                       linestyle='None',
                       marker='o',
                       markersize=15)
    fill, = plt.plot(df.loc[df['order_event_type'].isin([4])]['fill_price'],
                     color='red',
                     linestyle='None',
                     marker='s',
                     markersize=15)
    cancel, = plt.plot(df.loc[df['order_event_type'].isin([14])]['price'],
                       color='blue',
                       linestyle='None',
                       marker='x',
                       markersize=15)
    plt.legend([line, accept, fill, cancel], ['price', 'submit accept', 'fill', 'cancel'])
  elif exchange == 'Bitflyer':
    line, = plt.plot(df[df['price'].notnull()]['price'], color='yellow')
    accept, = plt.plot(df.loc[df['order_event_type'].isin([1, 2, 7])]['price'],
                       color='black',
                       linestyle='None',
                       marker='o',
                       markersize=15)
    fill, = plt.plot(df.loc[df['order_event_type'].isin([4])]['fill_price'],
                     color='red',
                     linestyle='None',
                     marker='s',
                     markersize=15)
    cancel, = plt.plot(df.loc[df['order_event_type'].isin([14])]['price'],
                       color='blue',
                       linestyle='None',
                       marker='x',
                       markersize=15)
    plt.legend([line, accept, fill, cancel], ['price', 'submit accept', 'fill', 'cancel'])
  else:
    raise ValueError(exchange)
  plt.savefig(os.path.join(csv_root, 'order_event_plot.png'))


def read_csv_into_df(csv_root, csv, timestamp_unit='ns'):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['timestamp'], unit=timestamp_unit)
  df = df.set_index('datetime')
  return df


def main(argv):
  csv_root = FLAGS.csv_root
  exchange = FLAGS.exchange
  assert csv_root, '--csv_root must be specified.'
  assert exchange, '--exchange must be specified.'

  balance_df = read_csv_into_df(csv_root, 'balance.csv')
  plot_balance(balance_df, csv_root)
  position_df = read_csv_into_df(csv_root, 'position.csv')
  plot_position(position_df, csv_root)
  order_event_df = read_csv_into_df(csv_root, 'order_event.csv', 'ns')
  plot_order_event(order_event_df, csv_root, exchange)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_root', None, 'Input csv files root directory.')

  flags.DEFINE_string('exchange', None, 'Exchange name.')

  app.run(main)
